# db.py
import psycopg2
from psycopg2 import extras
from .config import DB_CONFIG

def get_conn():
    return psycopg2.connect(**DB_CONFIG)

def fetch_all(query, params=None):
    conn = None
    try:
        conn = get_conn()
        cur = conn.cursor(cursor_factory=extras.RealDictCursor)
        cur.execute(query, params)
        rows = cur.fetchall()
        cur.close()
        return rows
    finally:
        if conn:
            conn.close()

def fetch_one(query, params=None):
    conn = None
    try:
        conn = get_conn()
        cur = conn.cursor(cursor_factory=extras.RealDictCursor)
        cur.execute(query, params)
        row = cur.fetchone()
        cur.close()
        return row
    finally:
        if conn:
            conn.close()

def execute(query, params=None, commit=True):
    conn = None
    try:
        conn = get_conn()
        cur = conn.cursor()
        cur.execute(query, params)
        if commit:
            conn.commit()
        cur.close()
    except Exception:
        if conn:
            conn.rollback()
        raise
    finally:
        if conn:
            conn.close()

# for multi-statement transactional operations: pass a function operating on a cursor
def transaction(fn):
    conn = None
    try:
        conn = get_conn()
        cur = conn.cursor(cursor_factory=extras.RealDictCursor)
        result = fn(cur)
        conn.commit()
        cur.close()
        return result
    except Exception:
        if conn:
            conn.rollback()
        raise
    finally:
        if conn:
            conn.close()
